import os
from os import getcwd
from xml.etree import ElementTree as ET
import re


# 定义一个创建一级分支object的函数
def create_object(root, xi, yi, xa, ya, obj_name):  # 参数依次为:树根，xmin, ymin, xmax, ymax
    # 创建异己分支object
    _object = ET.SubElement(root, 'object')
    # 创建二级分支
    name = ET.SubElement(_object, 'name')

    name.text = str(obj_name)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'

    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'

    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'

    # 创建bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % xi

    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % yi

    xmax = ET.SubElement(bndbox, 'xamx')
    xmax.text = '%s' % xa

    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % ya


# 定义一个创建xml文件的函数
def create_tree(images_path, images_name, h, w):
    global annotation
    # 创建树根annotation
    annotation = ET.Element('annotation')
    # 创建一级分支floder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder标签内容
    folder.text = images_path

    # 创建一级分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = images_name

    # 创建一级分支path
    path = ET.SubElement(annotation, 'path')
    path.text = getcwd() + '\{}\{}'.format(images_path, images_name)  # 用于返回当前工作目录

    # 创建一级分支source
    source = ET.SubElement(annotation, 'source')
    # 创建source下的二级分支database
    database = ET.SubElement(source, 'database')
    database.text = 'UnKnown'

    # 创建一级分支size
    size = ET.SubElement(annotation, 'size')
    # 创建size下的二级分支图像的宽、高级depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 创建一级分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'


if __name__ == '__main__':
    create_tree('data', '1', 200, 300)
    create_object(annotation, 100, 120, 150, 300, 'dog')

    # 将树模型写入xml文件
    tree = ET.ElementTree(annotation)
    # re.sub(r'(<[^/][^<>]*[^/]>)\s*([^<>]{,40}?)\s*(</[^<>]*>)', r'\1\2\3', tree)
    tree.write('.\{}\{}.xml'.format('data', '1'.strip('jpg')))
